import os
import re
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
import pandas as pd
from tqdm import tqdm

from LLM.Siliconflow import SiliconFlow
from main import main

from Hyper_Text.report_html import generate_medical_kg_report
from Hyper_Text.result_html import generate_medical_kg_result

def process_medical_document(doc_file: str):
    """
    处理单个病历文档，提取信息，调用LLM分析死亡因果链，并将结果保存为TXT文件。

    - 保留 tqdm 进度条（即使单行）
    - 输出文件命名: {原文件名}_{患者姓名}_processed_medical_data.txt
    - 内容为纯 JSON（自动去除可能的 ```json...``` 包裹）
    """
    base_name = os.path.splitext(os.path.basename(doc_file))[0]

    try:
        results = main(doc_file)
        patient_name = str(results.get('name', '未知患者')).strip() or '未知患者'
        admission = str(results.get('admission_diagnosis', '')).strip()
        death_record = str(results.get('death_record', '')).strip()
        death_discussion = str(results.get('death_discussion', '')).strip()
    except Exception as e:
        print(f"处理文档时发生错误: {e}")
        patient_name = '未知患者'
        admission = death_record = death_discussion = ''

    # 创建数据（保持 DataFrame 结构以兼容 progress_apply）
    data = {
        "姓名": [patient_name],
        "诊断摘要": [(admission, death_record, death_discussion)]
    }
    df = pd.DataFrame(data)

    # 定义Prompt模板
    template = """
你是一位专业的医学专家。请从提供的病历资料中分析死亡因果链并构建医学知识图谱。

**病历资料：**
入院诊断: {admission}
死亡记录: {death_record}
死亡讨论: {death_discussion}

**任务要求：**
1. **因果链分析**：识别并梳理死亡相关的疾病因果关系，构建三元组（前因疾病 → 关系 → 结果疾病）
2. **疾病编码**：将所有疾病标准化为ICD-10医学代码
3. **死亡链排序**：按照因果顺序排列疾病（1←2←3←4，数字越大为越根本的病因）
4. **分类标注**：区分直接死因、根本死因和促进因素

**输出格式：**
```json
{{
  "death_causality_chain": [
    {{
      "sequence": 1,
      "disease_name": "疾病名称",
      "icd10_code": "ICD-10代码",
      "role": "直接死因/根本死因/促进因素",
      "caused_by": "前一级疾病序号或null"
    }}
  ],
  "causality_triplets": [
    {{
      "cause": "前因疾病",
      "relation": "导致/加重/诱发",
      "effect": "结果疾病",
      "icd10_cause": "前因ICD-10",
      "icd10_effect": "结果ICD-10"
    }}
  ],
  "knowledge_graph_nodes": [
    {{
      "id": "节点ID",
      "label": "疾病名称",
      "icd10": "ICD-10代码",
      "type": "疾病类型"
    }}
  ],
  "knowledge_graph_edges": [
    {{
      "source": "源节点ID",
      "target": "目标节点ID",
      "relation": "关系类型"
    }}
  ]
}}  ```

仅输出JSON结构化数据，不添加任何解释文字。
"""

    prompt = PromptTemplate(
        input_variables=["admission", "death_record", "death_discussion"],
        template=template
    )

    llm = SiliconFlow()
    chain = LLMChain(llm=llm, prompt=prompt)

    def process_entry(entry):
        try:
            admission_, death_record_, death_discussion_ = entry
            print(f"处理数据: {admission_[:50]}...")
            result = chain.run({
                "admission": admission_,
                "death_record": death_record_,
                "death_discussion": death_discussion_
            }).strip()

            # 清理可能的 Markdown 代码块包裹
            result = re.sub(r'^```(?:json)?\s*', '', result)
            result = re.sub(r'\s*```$', '', result)
            return result
        except Exception as e:
            print(f"处理条目时出错: {e}")
            return f"ERROR: {str(e)}"

    # 使用 tqdm 显示进度条（保留你的要求）
    print("开始处理病历数据...")
    tqdm.pandas(desc="处理进度")
    df['processed'] = df['诊断摘要'].progress_apply(process_entry)

    # 检查结果
    print("\n处理完成！")
    success_count = df['processed'].notna().sum()
    fail_count = df['processed'].isna().sum()
    print(f"成功处理条目数: {success_count}")
    print(f"失败条目数: {fail_count}")

    # 生成安全的患者姓名（用于文件名）
    patient_name_clean = "".join(c for c in df['姓名'].iloc[0] if c.isalnum() or c in (' ', '_', '-')).strip()
    output_filename = f"{base_name}_{patient_name_clean}_processed_medical_data.txt"
    output_path = os.path.join(os.path.dirname(doc_file), output_filename)

    # 写入 TXT 文件（取第一个 processed 结果）
    output_content = df['processed'].iloc[0] if not df['processed'].isna().iloc[0] else "处理失败，无有效输出。"
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(output_content)
        print(f"结果已保存至: {output_path}")
    except Exception as e:
        print(f"保存 TXT 文件失败: {e}")

    return output_path, output_content  # 同时返回路径和内容（推荐）


def process_all_documents(dict_path):
    # 遍历目录中的所有文件
    for filename in os.listdir(dict_path):
        if filename.endswith(".docx"):
            doc_file = os.path.join(dict_path, filename)

            try:
                output_path, json_str = process_medical_document(doc_file)

                # 从 output_path 中提取基础文件名（不含扩展名）
                base_name = os.path.splitext(os.path.basename(output_path))[0]

                # 提取人名：匹配数字后、"_processed_medical_data"前的部分
                name_match = re.search(r'^\d+_(.*)_processed_medical_data$', base_name)
                if name_match:
                    person_name = name_match.group(1)
                else:
                    person_name = "unknown"

                # 构造报告前缀
                if base_name.endswith("_processed_medical_data"):
                    output_prefix = base_name[:-len("_processed_medical_data")] + "_report"
                else:
                    output_prefix = base_name + "_report"

                # 生成报告
                generate_medical_kg_report(json_str, output_prefix=output_prefix)
                generate_medical_kg_result(json_str, output_prefix=output_prefix, person_name=person_name)

                print(f"Processed: {filename} -> {output_prefix}")

            except Exception as e:
                print(f"Error processing {filename}: {e}")

if __name__ == "__main__":
    # dict_path = r'C:\Users\yyds\Desktop\new_demo'
    # doc_file = os.path.join(dict_path, "519038.docx")
    # output_path, json_str = process_medical_document(doc_file)
    #
    # # 从 output_path 中提取基础文件名（不含扩展名）
    # base_name = os.path.splitext(os.path.basename(output_path))[0]
    #
    # # 提取人名：匹配数字后、"_processed_medical_data"前的部分
    # name_match = re.search(r'^\d+_(.*)_processed_medical_data$', base_name)
    # if name_match:
    #     person_name = name_match.group(1)
    # else:
    #     # 如果格式不符合预期，尝试其他方式或设为默认
    #     person_name = "unknown"
    #
    # # 去掉末尾的 "_processed_medical_data" 部分
    # if base_name.endswith("_processed_medical_data"):
    #     output_prefix = base_name[:-len("_processed_medical_data")] + "_report"
    # else:
    #     # 如果格式不符合预期，也可以选择直接用整个 base_name + "_report"
    #     output_prefix = base_name + "_report"
    #
    # # 生成报告
    # generate_medical_kg_report(json_str, output_prefix=output_prefix)
    # generate_medical_kg_result(json_str, output_prefix=output_prefix, person_name=person_name)

    # 使用示例
    dict_path = r'C:\Users\yyds\Desktop\death_rerun'
    process_all_documents(dict_path)